import json
import os

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from pydantic.v1 import BaseModel

if __name__ == "__main__":
    def show_json(data):
        """用于展示json数据"""
        if isinstance(data, str):
            obj = json.loads(data)
            print(json.dumps(obj, indent=4))
        elif isinstance(data, dict) or isinstance(data, list):
            print(json.dumps(data, indent=4))
        elif issubclass(type(data), BaseModel):
            print(json.dumps(data.dict(), indent=4, ensure_ascii=False))


    def show_list_obj(data):
        """用于展示一组对象"""
        if isinstance(data, list):
            for item in data:
                show_json(item)
        else:
            raise ValueError("Input is not a list")

    os.environ["OPENAI_API_KEY"] = os.environ["OPENAI_API_KEY_ZHIHU"]
    os.environ["OPENAI_API_BASE"] = os.environ["OPENAI_API_BASE_ZHIHU"]
    os.environ["OPENAI_BASE_URL"] = os.environ["OPENAI_API_BASE_ZHIHU"]

    documents = SimpleDirectoryReader("./llamaIndexDatas",recursive=True,required_exts=[".pdf"]).load_data()
    index = VectorStoreIndex.from_documents(documents)

    vector_retriever = index.as_retriever(similarity_top_k=2)

    results = vector_retriever.retrieve("llama2有多少参数")
    show_list_obj(results)